#!/opt/local/bin/python

from analhighbeta import linacreflected
from scipy.optimize import fminbound, fmin_tnc
from rfvectors import cavity
import csv
from numpy import nan, array
import matplotlib.pyplot as plt

def lowbetareflected(splitinds, Ql, filename='lowbetanew.csv'):
    filehandles = []
    data = csv.reader(open(filename, 'rU'), dialect=csv.excel_tab, delimiter=',')
    rowcounter = 0
    secnum = 0
    for row in data:
        rowcounter += 1
        if rowcounter<splitinds[0]:
            filestuff(filehandles, row, 0)
        elif rowcounter<splitinds[1]:
            filestuff(filehandles, row, 1)
        else:
            filestuff(filehandles, row, 2)

    for i in filehandles:
        i.close()

    power = []
    for i in xrange(len(splitinds)+1):
        fname = 'section'+str(i)+'.csv'
        lowbetaargs  = (fname,  704.42e6, nan, 50e-3, 50e-3, 300.0, 0.70, nan, False)
        Qlopt = fminbound(linacreflected, 1e4, 1e7, args=lowbetaargs)
        power.append(linacreflected(Ql=Qlopt, filename=fname, f=704.42e6, delF=nan, 
            I_b0=50e-3, I_b=50e-3, RoverQ=300.0, beta_0=0.7, Eampfile=nan))
    return sum(power)

def filestuff(fhandlist, dat, indnum):
    try:
        fhandlist[indnum].write(','.join(dat)+'\n')
    except IndexError:
        fhandlist.append(open('section'+str(indnum)+'.csv', 'w'))
        fhandlist[indnum].write(','.join(dat)+'\n')

def minhelper(x0):
    #print "%0.0f, %0.0f, %0.3e, %0.3e, %0.3e" % (x0[0],x0[0]+x0[1],x0[2],x0[3],x0[4])
    return lowbetareflected(splitinds=[x0[0], x0[0]+x0[1]], Ql=[1e5,1e5,1e5])

if __name__=="__main__":
    totref = lowbetareflected(splitinds=[6,16], Ql=[1e6,5e5,1e5])
    print "Total reflected = %f" % totref

    refl = []
    inputrange = range(2,35)
    for start in inputrange:
        res, nfeval, rc = fmin_tnc(minhelper, x0=[start,7], approx_grad=True, epsilon=4.1,
                bounds=[(2,30), (2,30)],
                maxfun=1000)
        print "nfeval = %d :: return code = %d" % (nfeval, rc)
        print "split locations = %0.0f, %0.0f" % (36+res[0], 36+res[0]+res[1])
        #print "Ql = (%0.3e, %0.3e, %0.3e)" % (res[2], res[3], res[4])
        print "start = %f" % start
        print "Total reflected after optimisation = %0.3f kW" % (minhelper(res)/1e3)
        refl.append(minhelper(res))

    fig = plt.figure(1)
    ax = fig.add_subplot(111)
    ax.plot(inputrange, array(refl), '-x')
    ax.grid()
    plt.show()

